#! /usr/bin/env python
# -*- coding: utf-8 -*-

# file: get_specific_list.py
# author: jiangqr
# describe: filter list file about voc or imagenet images
# time: 2016-11-17

import os
import xml.dom.minidom as minidom
import caffe
from caffe.model_libs import *
from google.protobuf import text_format

import cv2


def get_object_id_from_xml(file):
    dom = minidom.parse(file)
    names = dom.getElementsByTagName("name")
    if len(names) == 1:
        return names[0].childNodes[0].data
    else:
        return 0


def get_file_list(directory):
    xml_file_list = [];
    file_names = os.listdir(directory)
    for file_name in file_names:
        file_path = os.path.join('%s%s' % (directory, file_name))
        xml_file_list.append(file_path)
    return xml_file_list

# main

rootDir = '/home/ss/data/ILSVRC/'
labelMapFile = '/home/ss/Git/caffe/data/ILSVRC2016/labelmap_ilsvrc_det.100class.prototxt'
oriListFile = '/home/ss/Git/caffe/data/ILSVRC2016/val2.txt'
newListFile = '/home/ss/Git/caffe/data/ILSVRC2016/val2_100class.txt'
nameSizeiFile = '/home/ss/Git/caffe/data/ILSVRC2016/val2_100class_name_size.txt'
bSaveSizeFile = 1

labelmapFileBuf = open(labelMapFile).read()
label_map = caffe_pb2.LabelMap()
text_format.Merge(labelmapFileBuf, label_map)
name_label = label_map.item

name_list = []
for name in name_label:
    name_list.append(name.name)

save_file_list = []
save_size_file_list = []
i = 0
sizeFileIndex = 1
with open(oriListFile) as f:
    while 1:
        print i
        i += 1
        line = f.readline()
        if not line:
            break
        imgPath, xmlPath = line.split(' ')
        xmlPath = xmlPath.strip()
        imgPath = imgPath.strip()

        name = get_object_id_from_xml(rootDir+xmlPath)
        if name in name_list:
            save_file_list.append(line)

            if bSaveSizeFile:
                img = cv2.imread(rootDir + imgPath)
                height, width, channels = img.shape
                save_size_file_list.append('{} {} {}'.format(sizeFileIndex, height, width))
                sizeFileIndex += 1


with open(newListFile,'wb') as f:
    for path in save_file_list:
        f.write(path)

if bSaveSizeFile:
    with open(nameSizeiFile, 'wb') as f:
        for list in save_size_file_list:
            f.write(list)
            f.write('\n')
